/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "parvec_axpby.h"
#include "parvec_axpbypcz.h"
#include "par_chebyshev_polynomial.h"

namespace YHAMG
{

Par_Chebyshev_Polynomial::Par_Chebyshev_Polynomial(const Par_Operator& _A, double _alpha, double _beta, int _degree)
	: Par_Operator(_A.comm),
	A(&_A),
	alpha(_alpha),
	beta(_beta),
	degree(_degree)
{
}

int Par_Chebyshev_Polynomial::InSize() const
{
	return A->InSize();
}

int Par_Chebyshev_Polynomial::OutSize() const
{
	return A->OutSize();
}

void Par_Chebyshev_Polynomial::Apply(const Par_Vector& x, const Par_Vector& y) const
{
	double theta, delta, sigma, rho, rho1;
	Par_Vector z(comm), w(comm);
	z.Resize(OutSize());
	w.Resize(OutSize());

	theta = (beta + alpha) * 0.5;
	delta = 2.0 / (beta - alpha);
	sigma = theta * delta;
	rho = 1.0 / sigma;

	y.Fill(0.0);
	ParVec_AXPBY(0.0, w, 1.0 / theta, x, w);

	for (int k = 0; k < degree; ++k)
	{
		rho1 = 1.0 / (2.0 * sigma - rho);
		y.AddScaled(1.0, w);
		A->Apply(y, z);
		ParVec_AXPBYPCZ(rho * rho1, w, 2.0 * rho1 * delta, x, -2.0 * rho1 * delta, z, w);
		rho = rho1;
	}
}

void Par_Chebyshev_Polynomial::Apply(const Par_MultiVector& X, const Par_MultiVector& Y) const
{
	int m = X.num_vectors;
	double theta, delta, sigma, rho, rho1;
	Par_MultiVector Z(comm), W(comm);
	Z.Allocate(OutSize(), m);
	W.Allocate(OutSize(), m);

	theta = (beta + alpha) * 0.5;
	delta = 2.0 / (beta - alpha);
	sigma = theta * delta;
	rho = 1.0 / sigma;

	for (int j = 0; j < m; ++j)
	{
		Y(j).Fill(0.0);
		ParVec_AXPBY(0.0, W(j), 1.0 / theta, X(j), W(j));
	}

	for (int k = 0; k < degree; ++k)
	{
		rho1 = 1.0 / (2.0 * sigma - rho);	
		for (int j = 0; j < m; ++j)
			Y(j).AddScaled(1.0, W(j));
		A->Apply(Y, Z);
		for (int j = 0; j < m; ++j)
			ParVec_AXPBYPCZ(rho * rho1, W(j), 2.0 * rho1 * delta, X(j), -2.0 * rho1 * delta, Z(j), W(j));
		rho = rho1;
	}
}

}